#!/usr/bin/env python

from __future__ import with_statement
from PIL import Image, ImageEnhance, ImageOps
from Crypto.Cipher import AES as cipher

##### BLUR #####
def blur(pil_image, factor = 1.0):
    return ImageEnhance.Sharpness(pil_image).enhance(1.0 - factor)
    
##### SHARPEN #####
def sharpen(pil_image, factor = 1.0):
    return ImageEnhance.Sharpness(pil_image).enhance(factor + 1.0)

##### BRIGHTNESS #####
def brightness(pil_image, factor = 0.75):
    return ImageEnhance.Brightness(pil_image).enhance(factor)

##### COLOR #####
def color(pil_image, factor = 0.75):
    return ImageEnhance.Color(pil_image).enhance(factor)

##### CONTRAST #####
def contrast(pil_image, factor = 0.75):
    return ImageEnhance.Contrast(pil_image).enhance(factor)

##### INVERT #####
def invert(pil_image):
    return ImageOps.invert(pil_image)

##### FLIP HORIZONTAL #####
def flip_horizontal(pil_image):
    return pil_image.transpose(Image.FLIP_TOP_BOTTOM)
    
##### FLIP VERTICAL #####
def flip_vertical(pil_image):
    return pil_image.transpose(Image.FLIP_LEFT_RIGHT)

##### ROTATE #####
def rotate(pil_image, angle = 45):
    return pil_image.rotate(int(angle))

##### SCALE #####
def scale(pil_image, size = (400, 300), maintain_aspect = True, method = 1): # 1 = Image.ANTIALIAS
    # This function calculates the aspect maintaining size that the image
    # should be resized to.
    def adjust_size_to_maintain_aspect(original_size, desired_size):
        scale_ = float(desired_size[0])/original_size[0]
        if original_size[1] * scale_ <= desired_size[1]:
            height = int(original_size[1] * scale_)
            if height > desired_size[1]:
                height = desired_size[1]
            return (desired_size[0], height)
        else:
            scale_ = float(desired_size[1])/original_size[1]
            width = int(original_size[0] * scale_)
            if width > desired_size[0]:
                width = desired_size[0]
            return (width, desired_size[1])
    
    if maintain_aspect:
        size = adjust_size_to_maintain_aspect(pil_image.size, size)
    
    # Scale the image.
    return pil_image.resize(size, method)

##### Encrypt (AES) #####
def encrypt_aes(data):
    obj = cipher.new('very secret key!')
    data_size = len(data)
    padding = 0 if data_size % 16 == 0 else (16 - data_size%16)
    if padding != 0:
        return obj.encrypt(data + 'X'*padding)
    else:
        return obj.encrypt(data)


##### Mandelbrot fractal renderer #####
def mandelbrot(start, end, width, height, iterations):
    data = []
    # Create the fractal.
    hstep = (end.real-start.real)/width
    vstep = (end.imag-start.imag)/height
    cstep = 255.0/iterations
    black = chr(255)
    for i in range (0, height):
        current_vstep = start.imag + vstep*i
        current_hstep = 0.0
        for j in range (0, width):
            z = (0+0j)
            c = complex(start.real + current_hstep, current_vstep)
            current_hstep += hstep
            color = 0.0
            k = 0
            while k < iterations:
                z *= z
                z += c
                if abs(z) > 2:
                    # We have reached the limit. Write a pixel of the current color.
                    outcol = int(color)
                    data.append(chr(outcol))
                    break
                color += cstep
                k += 1
            # If nothing was written in the loop we write a black pixel.
            if k == iterations:
                data.append(black)
    return ''.join(data)


# Main code below...
import sys
from time import time

# Check arguments.
if len(sys.argv) < 3:
    print "Usage: measure_weight.py input-image host-strength [-i iterations]"
    sys.exit(1)

# Read the input image.
image = None
try:
    with open(sys.argv[1], 'r') as infile:
        image_data = infile.read()
    pil_image = Image.open(sys.argv[1])
except:
    print 'Error reading image.'
    sys.exit(1)

# Read the host strength.
strength = 0
try:
    strength = float(sys.argv[2])
except:
    print 'Error reading host strength.'
    sys.exit(1)

# Check for an optional iterations count.
iterations = 10
if '-i' in sys.argv:
    try:
        iterations = int(sys.argv[sys.argv.index('-i') + 1])
    except:
        print 'Error parsing number of iterations.'
        sys.exit(1)

# Start measuring the weight of all of the image operations.
image_functions = [('brightness', brightness), ('color', color), ('contrast', contrast), ('sharpen', sharpen),
                               ('blur', blur), ('scale', scale), ('invert', invert), ('flip_horizontal', flip_horizontal),
                               ('flip_vertical', flip_vertical), ('rotate', rotate)]

# Image functions
for name, function in image_functions:
    total = 0.0
    for _ in range(0, iterations):
        start = time()
        function(pil_image)
        stop = time()
        total += stop-start
    print name, "%.3f"%((total/iterations)*strength)

# Encryption
total = 0.0
for _ in range(0, iterations):
    start = time()
    encrypt_aes(image_data)
    stop = time()
    total += stop-start
print 'encrypt', "%.3f"%((total/iterations)*strength)

# Mandelbrot
total = 0.0
for _ in range(0, iterations):
    start = time()
    mandelbrot((-1.5+-1j), (0.5+1j), 640, 480, 25)
    stop = time()
    total += stop-start
print 'mandelbrot', "%.3f"%((total/iterations)*strength)
